# -*- coding: utf-8 -*-
from Reader import Reader
import matplotlib.pyplot as plt
import matplotlib
matplotlib.rcParams['font.family'] = ['KaiTi', 'sans-serif']
matplotlib.rcParams['axes.unicode_minus'] = False
plt.style.use("./journal.mplstyle")
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.ticker import MultipleLocator
import config
import numpy as np
import h5py
import statsmodels.api as sm
reader = Reader('data.xlsx')
time_12_5 = np.arange(6)
time_13_5 = np.arange(6,11)
industryName = config.industryName[np.array([0,2,3,4,5])]
years = reader.T

# gamma
gammaselect = np.array([0, 1,2,3])
gamma_res_12_5 = np.zeros((5,2), dtype=np.float64)# b k
gamma_res_13_5 = np.zeros((5,2), dtype=np.float64)
GDPs = reader.P_GDP[['GDP_'+i for i in industryName[gammaselect]]].to_numpy()
GDPs_norm = GDPs/(np.sum(GDPs, axis=1)[:, np.newaxis])
selyears = years[time_12_5]
X = sm.add_constant(selyears-selyears[0])
for i in gammaselect[:-1]:
    model = sm.OLS(GDPs_norm[time_12_5][:,i],X)
    results = model.fit()
    gamma_res_12_5[i] = results.params
gamma_res_12_5[3] = [1-np.sum(gamma_res_12_5[gammaselect, 0]), -np.sum(gamma_res_12_5[gammaselect, 1])]
selyears = years[time_13_5]
X = sm.add_constant(selyears-selyears[0])
for i in gammaselect[:-1]:
    model = sm.OLS(GDPs_norm[time_13_5][:,i],X)
    results = model.fit()
    gamma_res_13_5[i] = results.params
gamma_res_13_5[3] = [1-np.sum(gamma_res_13_5[gammaselect, 0]), -np.sum(gamma_res_13_5[gammaselect, 1])]
print(gamma_res_12_5, gamma_res_13_5)
# beta
beta_res_12_5 = np.zeros((5,2), dtype=np.float64)# b k
beta_res_13_5 = np.zeros((5,2), dtype=np.float64)
Energys = []
for industry in industryName:
    targetName = [industry + '.' +i for i in config.energyName]
    Energys.append(np.sum(reader.Energy[targetName].to_numpy(), axis=1))
Energys = np.array(Energys).T
GDPs_t = np.hstack([GDPs, reader.P_GDP['GDP_t'].to_numpy()[:,np.newaxis]])
Energy_GDPs = Energys/GDPs_t
selyears = years[time_12_5]
X = sm.add_constant(selyears-selyears[0])
for i in range(len(industryName)):
    model = sm.OLS(Energy_GDPs[time_12_5][:,i], X)
    results = model.fit()
    beta_res_12_5[i] = results.params
selyears = years[time_13_5]
X = sm.add_constant(selyears-selyears[0])
for i in range(len(industryName)):
    model = sm.OLS(Energy_GDPs[time_13_5][:,i], X)
    results = model.fit()
    beta_res_13_5[i] = results.params
print(beta_res_12_5, beta_res_13_5)
# zeta
zeta_res_12_5 = np.zeros((5,6,2), dtype=np.float64)# b k
zeta_res_13_5 = np.zeros((5,6,2), dtype=np.float64)
selyears1 = years[time_12_5]
selyears2 = years[time_13_5]
X1 = sm.add_constant(selyears1-selyears1[0])
X2 = sm.add_constant(selyears2-selyears2[0])
for i in range(len(industryName)):
    ys = []
    for j in range(len(config.energyName)):
        targetName = industryName[i] + '.' +config.energyName[j]
        ys.append(reader.Energy[targetName])
    ys = np.array(ys)
    ys_sum = np.sum(ys, axis=0)
    for j in range(len(config.energyName)-1):
        targetName = industryName[i] + '.' +config.energyName[j]
        ys = reader.Energy[targetName]
        model = sm.OLS(ys[time_12_5]/ys_sum[time_12_5], X1)
        results = model.fit()
        zeta_res_12_5[i,j] = results.params
        model = sm.OLS(ys[time_13_5]/ys_sum[time_13_5], X2)
        results = model.fit()
        zeta_res_13_5[i,j] = results.params
    zeta_res_12_5[i,-1] = [1-np.sum(zeta_res_13_5[i, :-1, 0]), -np.sum(zeta_res_13_5[i, :-1, 1])]
print(zeta_res_12_5, zeta_res_13_5)
# alpha
alpha_res_12_5 = np.zeros((5,6,2), dtype=np.float64)# b k
alpha_res_13_5 = np.zeros((5,6,2), dtype=np.float64)
for i in range(len(industryName)):
    for j in range(len(config.energyName)-1):
        targetName = industryName[i] + '.' +config.energyName[j]
        ys = reader.Carbon_Energy[targetName].to_numpy()
        if not np.all(ys[time_12_5]=='-'):
            alpha_res_12_5[i, j, 0] = np.mean(ys[time_12_5][ys[time_12_5]!='-'])#results.params
        if not np.all(ys[time_13_5]=='-'):
            alpha_res_13_5[i,j, 0] = np.mean(ys[time_13_5][ys[time_13_5]!='-'])#results.params
            print('{}:{},{}'.format(targetName, alpha_res_13_5[i,j, 0], ys[time_13_5][ys[time_13_5]!='-']))
print(alpha_res_12_5, alpha_res_13_5)
# G/P P
GDP_t, P = reader.P_GDP['GDP_t'].to_numpy(), reader.P_GDP['P'].to_numpy()
GDP_P = GDP_t/P
GP_res = np.zeros((2,2), dtype=np.float64)
X = sm.add_constant(years-years[0])
model = sm.OLS(GDP_P, X)
results = model.fit()
GP_res[0, :2] = results.params
X = sm.add_constant(np.exp(-(years[5:]-years[0])))
model = sm.OLS(P[5:], X)
results = model.fit()
GP_res[1, :2] = results.params

with h5py.File('model.h5', 'w') as opt:
    opt.create_dataset('12_5/gamma', data=gamma_res_12_5, compression='gzip')
    opt.create_dataset('13_5/gamma', data=gamma_res_13_5, compression='gzip')
    opt.create_dataset('12_5/beta', data=beta_res_12_5, compression='gzip')
    opt.create_dataset('13_5/beta', data=beta_res_13_5, compression='gzip')
    opt.create_dataset('12_5/zeta', data=zeta_res_12_5, compression='gzip')
    opt.create_dataset('13_5/zeta', data=zeta_res_13_5, compression='gzip')
    opt.create_dataset('12_5/alpha', data=alpha_res_12_5, compression='gzip')
    opt.create_dataset('13_5/alpha', data=alpha_res_13_5, compression='gzip')
    opt.create_dataset('GP', data=GP_res, compression='gzip')

with PdfPages('model.h5.pdf') as pdf:
    fig, ax = plt.subplots()
    ax.plot(years[time_12_5], (years[time_12_5]-years[time_12_5][0])[:,np.newaxis]*gamma_res_12_5[:,1][np.newaxis, :]+gamma_res_12_5[:,0])
    ax.plot(years[time_13_5], (years[time_13_5]-years[time_13_5][0])[:,np.newaxis]*gamma_res_13_5[:,1][np.newaxis, :]+gamma_res_13_5[:,0])
    ax.set_xlabel('年份')
    ax.xaxis.set_major_locator(MultipleLocator(1))
    ax.set_ylabel('不同产业部门GDP比例')
    pdf.savefig(fig)
    plt.close()

    fig, ax = plt.subplots()
    ax.plot(years[time_12_5], (years[time_12_5]-years[time_12_5][0])[:,np.newaxis]*beta_res_12_5[:,1][np.newaxis, :]+beta_res_12_5[:,0])
    ax.plot(years[time_13_5], (years[time_13_5]-years[time_13_5][0])[:,np.newaxis]*beta_res_13_5[:,1][np.newaxis, :]+beta_res_13_5[:,0], marker='o')
    ax.set_xlabel('年份')
    ax.xaxis.set_major_locator(MultipleLocator(1))
    ax.set_ylabel('单位生产总值对应的能源消耗')
    pdf.savefig(fig)
    ax.set_yscale('log')
    pdf.savefig(fig)
    plt.close()

    for i in range(len(industryName)):
        fig, ax = plt.subplots()
        ax.stackplot(years[time_12_5], ((years[time_12_5]-years[time_12_5][0])[:,np.newaxis]*zeta_res_12_5[i,:,1][np.newaxis, :]+zeta_res_12_5[i,:,0]).T)
        ax.stackplot(years[time_13_5], ((years[time_13_5]-years[time_13_5][0])[:,np.newaxis]*zeta_res_13_5[i,:,1][np.newaxis, :]+zeta_res_13_5[i,:,0]).T)
        ax.set_xlabel('年份')
        ax.xaxis.set_major_locator(MultipleLocator(1))
        ax.set_ylabel('能源结构')
        pdf.savefig(fig)
        plt.close()
    
    for i in range(len(industryName)):
        fig, ax = plt.subplots()
        ax.plot(years[time_12_5], ((years[time_12_5]-years[time_12_5][0])[:,np.newaxis]*alpha_res_12_5[i,:,1][np.newaxis, :]+alpha_res_12_5[i,:,0]))
        ax.plot(years[time_13_5], ((years[time_13_5]-years[time_13_5][0])[:,np.newaxis]*alpha_res_13_5[i,:,1][np.newaxis, :]+alpha_res_13_5[i,:,0]))
        ax.set_xlabel('年份')
        ax.xaxis.set_major_locator(MultipleLocator(1))
        ax.set_ylabel('碳转换因子')
        pdf.savefig(fig)
        plt.close()

    fig, ax = plt.subplots()
    for i, ene in enumerate(config.energyName):
        ax.scatter([*[config.industryNameMap[i]+'_12' for i in industryName], *[config.industryNameMap[i]+'_13' for i in industryName]], [*alpha_res_12_5[:,i,0], *alpha_res_13_5[:,i,0]], marker='o', label=config.energyNameMap[ene], )
        print([*alpha_res_12_5[:,i,0], *alpha_res_13_5[:,i,0]])
    ax.xaxis.set_major_locator(MultipleLocator(1))
    ax.set_ylabel('碳转换因子')
    ax.tick_params(axis='x', rotation=90)
    ax.legend()
    pdf.savefig(fig)
    plt.close()

    fig, ax = plt.subplots(figsize=(12, 6))
    ax.scatter(years, GDP_P, marker='o', label='人均GDP数据')
    years_predict = [*years,*(years+11)]
    ax.plot(years_predict, (years_predict-years_predict[0])*GP_res[0,1]+GP_res[0,0], color='r', label='人均GDP预测')
    ax.xaxis.set_major_locator(MultipleLocator(1))
    ax.set_xlabel('年份')
    ax.set_ylabel('人均GDP')
    ax.legend()
    pdf.savefig(fig)
    plt.close()

    fig, ax = plt.subplots(figsize=(12, 6))
    ax.scatter(years, P, marker='o', label='总人口数据')
    ax.plot(years_predict[5:], np.exp(-(years_predict[5:]-years_predict[0]))*GP_res[1,1]+GP_res[1,0], color='r', label='总人口预测')
    ax.xaxis.set_major_locator(MultipleLocator(1))
    ax.set_xlabel('年份')
    ax.set_ylabel('总人口')
    ax.legend()
    pdf.savefig(fig)
    plt.close()